import torch
from transformers import AutoModelForImageClassification, AutoTokenizer, pipeline
from PIL import Image

class MMRoleLite:
    def __init__(self, config):
        self.config = config
        self.setup_models()
        
    def setup_models(self):
        """初始化模型"""
        try:
            # 初始化视觉模型
            self.vision_model = pipeline(
                "image-classification",
                model=self.config.MODEL_CONFIG["vision_model"],
                device=self.config.MODEL_CONFIG["device"]
            )
            
            # 初始化文本生成模型
            self.text_generator = pipeline(
                "text-generation",
                model=self.config.MODEL_CONFIG["text_model"],
                device=self.config.MODEL_CONFIG["device"]
            )
            print("模型加载成功！")
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            raise
        
    def analyze_image(self, image_path):
        """简化版图像分析"""
        try:
            image = Image.open(image_path)
            # 获取图像分类结果
            results = self.vision_model(image)
            # 返回最可能的类别描述
            return results[0]['label'] if results else "无法识别的图像"
        except Exception as e:
            print(f"图像分析错误: {str(e)}")
            return "图像处理出错"
    
    def generate_response(self, character_id, image_desc, user_input):
        """生成角色回应"""
        try:
            character = self.config.CHARACTERS[character_id]
            context = f"""
            作为{character['name']}，具有以下特征：
            性格: {character['personality']}
            背景: {character['background']}
            说话风格: {character['style']}
            
            看到一张图片显示: {image_desc}
            
            回应问题: {user_input}
            """
            
            response = self.text_generator(
                context,
                max_length=self.config.MODEL_CONFIG["max_length"],
                num_return_sequences=1
            )
            return response[0]['generated_text']
        except Exception as e:
            print(f"响应生成错误: {str(e)}")
            return "生成回应时出错"
    
    def evaluate_response(self, response, character_id, image_desc):
        """简化版评估"""
        try:
            scores = {}
            for dimension, metrics in self.config.EVAL_METRICS.items():
                # 简单评分示例
                score = len(response) / self.config.MODEL_CONFIG["max_length"]
                scores[dimension] = min(max(score * 10, 0), 10)  # 归一化到0-10
            return scores
        except Exception as e:
            print(f"评估错误: {str(e)}")
            return {"error": "评估过程出错"} 